from numpy import pi, cos, sin, empty, linspace, random, ogrid
from tvtk.api import tvtk
from mayavi import mlab
from pdb import set_trace

def coordinates(R=0.5, Z=3., dims=(10,10,30)):
    """ Define the cylinder geometry, and create the array of the coordinates
        corresponding to the geometry.
        R : cylinder radius
        Z : cylinder height
        dims : tuple (3) of the number of points along the r, q, z coordinates
        return numpy array of dimension dims x 3
    """
    nr, nq, nz = [d*1j for d in dims]
    r, q, z = ogrid[0:R:nr, 0:1.7*pi:nq, 0:Z:nz]

    x = cos(q)*r
    y = sin(q)*r

    pts = empty((x.shape[0], y.shape[1], z.shape[2], 3), dtype=float)
    pts[..., 0] = x
    pts[..., 1] = y
    pts[..., 2] = z

    return pts

def temperature(pts):
    """
    Define the temperature over the cylinder.
    pts : numpy [x1, x2, x3, 3] array of points
    """
    T = pts[...,0]**2 + pts[...,1]**2 + 0.3 - 0.1*pts[...,2]**2 / 4.
    return T

def set_data():
    R = 0.5
    Z = 3.
    dims = (10, 10, 30)

    pts = coordinates(R=R, Z=Z, dims=dims)
    T = temperature(pts)

    pts = pts.transpose(2, 1, 0, 3).copy()
    pts.shape = pts.size / 3, 3
    T = T.T.copy().ravel()

    sg = tvtk.StructuredGrid(dimensions=dims, points=pts)
    sg.point_data.scalars = T
    sg.point_data.scalars.name = 'temperature'

    return sg

def make_plot(sg):
    """
    Make a plot of the vtk data
    vtk_data : vtk data
    """
    d = mlab.pipeline.add_dataset(sg)
    surf = mlab.pipeline.surface(d, opacity=0.6)
    iso_surface = mlab.pipeline.iso_surface(d, contours=[0.1,0.2,0.3], opacity=0.2)
    iso_surface.actor.property.representation = 'wireframe'
    e = mlab.get_engine()
    s = e.scenes[0]
    s.scene.background = (1.0, 1.0, 1.0)
    s.scene.x_minus_view()
    mlab.show()

if __name__ == '__main__':
    from mayavi.scripts import mayavi2

    @mayavi2.standalone
    def view(sg):
        """
        View the data vtk_data in the Mayavi application.
        vtk_data : vtk data
        """
        d = mlab.pipeline.add_dataset(sg)
        surf = mlab.pipeline.surface(d, opacity=0.6)
        iso_surface = mlab.pipeline.iso_surface(d, contours=[0.1,0.2,0.3], opacity=0.2)
        iso_surface.actor.property.representation = 'wireframe'

    sg = set_data()
    view(sg)
